#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
OCR引擎模块
封装PaddleOCR功能
"""

# 导入PaddleOCR，但使用异常处理来应对可能的兼容性问题
PADDLE_OCR_AVAILABLE = False
PaddleOCRClass = None

try:
    from paddleocr import PaddleOCR
    PaddleOCRClass = PaddleOCR
    PADDLE_OCR_AVAILABLE = True
    print("成功导入PaddleOCR主类")
except ImportError:
    print("警告: 无法导入paddleocr.PaddleOCR，请检查安装")
except Exception as e:
    print(f"警告: PaddleOCR导入失败: {str(e)}")

import cv2
import numpy as np
import os

class OCREngine:
    def __init__(self):
        """
        初始化PaddleOCR引擎
        """
        if not PADDLE_OCR_AVAILABLE or PaddleOCRClass is None:
            self.model = None
            return
            
        # 初始化PaddleOCR模型
        try:
            # 使用标准方式初始化PaddleOCR
            #self.model = PaddleOCRClass(
            #    use_doc_orientation_classify=False, 
            #    use_doc_unwarping=False, 
            #    use_textline_orientation=False)
            # 使用PPOCRV5 mobile模型初始化PaddleOCR提升速度
            self.model = PaddleOCRClass(
                text_detection_model_name="PP-OCRv5_mobile_det",
                text_recognition_model_name="PP-OCRv5_mobile_rec",
                use_doc_orientation_classify=False,
                use_doc_unwarping=False,
                use_textline_orientation=False)
            print("PaddleOCR初始化成功")
        except Exception as e:
            print(f"PaddleOCR初始化失败: {str(e)}")
            self.model = None
            # 不抛出异常，而是设置model为None，让程序可以继续运行
    
    def recognize(self, image_path):
        """
        对图像进行OCR识别
        
        Args:
            image_path (str): 图像文件路径
            
        Returns:
            dict: 识别结果，包含文本和位置信息
        """
        if not PADDLE_OCR_AVAILABLE or self.model is None:
            return {"error": "错误: PaddleOCR不可用"}
            
        try:
            # 使用PaddleOCR进行识别
            result = self.model.predict(image_path)
            
            # 提取文本内容
            text_results = []
            if result is not None:
                # 根据1_res.json的数据结构和2.py的示例，处理PaddleOCR 3.2.0返回的结果格式
                # 结果是一个包含OCRResult对象的列表
                for res in result:
                    # 检查是否有rec_texts属性（包含识别的文本）
                    if hasattr(res, 'rec_texts'):
                        # 直接使用识别的文本列表
                        text_results.extend(res.rec_texts)
                    # 如果res本身就有text属性（兼容不同版本）
                    elif hasattr(res, 'text'):
                        text_results.append(res.text)
                    # 如果res是字典格式（兼容不同版本）
                    elif isinstance(res, dict) and 'rec_texts' in res:
                        text_results.extend(res['rec_texts'])
            
            return {"text_results": text_results, "raw_output": result}
        except Exception as e:
            error_msg = f"OCR识别过程中发生错误: {str(e)}"
            print(error_msg)
            return {"error": error_msg}
    
    def recognize_from_array(self, image_array):
        """
        对图像数组进行OCR识别
        
        Args:
            image_array (numpy.ndarray): 图像数组
            
        Returns:
            dict: 识别结果，包含文本和位置信息
        """
        if not PADDLE_OCR_AVAILABLE or self.model is None:
            return {"error": "错误: PaddleOCR不可用"}
            
        try:
            # 使用PaddleOCR进行识别
            result = self.model.predict(image_array)
            
            # 提取文本内容
            text_results = []
            if result is not None:
                # 根据1_res.json的数据结构和2.py的示例，处理PaddleOCR 3.2.0返回的结果格式
                # 结果是一个包含OCRResult对象的列表
                for res in result:
                    # 检查是否有rec_texts属性（包含识别的文本）
                    if hasattr(res, 'rec_texts'):
                        # 直接使用识别的文本列表
                        text_results.extend(res.rec_texts)
                    # 如果res本身就有text属性（兼容不同版本）
                    elif hasattr(res, 'text'):
                        text_results.append(res.text)
                    # 如果res是字典格式（兼容不同版本）
                    elif isinstance(res, dict) and 'rec_texts' in res:
                        text_results.extend(res['rec_texts'])
            
            return {"text_results": text_results, "raw_output": result}
        except Exception as e:
            error_msg = f"OCR识别过程中发生错误: {str(e)}"
            print(error_msg)
            return {"error": error_msg}
    
    def get_detailed_results(self, image_path):
        """
        获取详细的OCR识别结果，包括文本和位置信息
        
        Args:
            image_path (str): 图像文件路径
            
        Returns:
            list: 详细的识别结果列表
        """
        if not PADDLE_OCR_AVAILABLE or self.model is None:
            return []
            
        try:
            # 使用PaddleOCR进行识别
            result = self.model.predict(image_path)
            
            # 提取详细信息
            detailed_results = []
            if result is not None:
                # 根据1_res.json的数据结构和2.py的示例，处理PaddleOCR 3.2.0返回的结果格式
                # 结果是一个包含OCRResult对象的列表
                for res in result:
                    # 检查是否有rec_texts、rec_scores和rec_boxes属性
                    if hasattr(res, 'rec_texts') and hasattr(res, 'rec_scores') and hasattr(res, 'rec_boxes'):
                        # 使用OCRResult对象的属性，按照1_res.json的数据结构
                        for i in range(len(res.rec_texts)):
                            detailed_results.append({
                                "text": res.rec_texts[i],
                                "confidence": float(res.rec_scores[i]) if i < len(res.rec_scores) else 0.0,
                                "coordinates": res.rec_boxes[i].tolist() if i < len(res.rec_boxes) else []
                            })
                    # 如果res是字典格式（兼容不同版本）
                    elif isinstance(res, dict) and 'rec_texts' in res and 'rec_scores' in res and 'rec_boxes' in res:
                        for i in range(len(res['rec_texts'])):
                            detailed_results.append({
                                "text": res['rec_texts'][i],
                                "confidence": float(res['rec_scores'][i]) if i < len(res['rec_scores']) else 0.0,
                                "coordinates": res['rec_boxes'][i].tolist() if i < len(res['rec_boxes']) else []
                            })
            
            return detailed_results
        except Exception as e:
            print(f"获取详细OCR结果时发生错误: {str(e)}")
            return []